# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD 3-Clause license found in the
# LICENSE file in the root directory of this source tree.

from dataclasses import dataclass
from itertools import product
from typing import List, Tuple

import torch
from tabulate import tabulate
from tqdm import tqdm

# Assuming these imports based on the kernel location
from benchmarks.utils import benchmark_cuda_function_in_microseconds
from torchao.prototype.blockwise_fp8_training.kernels import (
    triton_fp8_blockwise_weight_quant_rhs,
)

device = torch.device("cuda")

# Needed since changing args to function causes recompiles
torch._dynamo.config.cache_size_limit = 1000


@dataclass(frozen=True)
class ExperimentConfig:
    input_shape: Tuple[int, int]  # (M, N)
    block_size: int


@dataclass(frozen=True)
class ExperimentResult:
    # time
    torch_us: float
    triton_us: float
    # mem bw
    torch_gbps: float
    triton_gbps: float


@dataclass(frozen=True)
class Experiment:
    config: ExperimentConfig
    result: ExperimentResult


def get_configs() -> List[ExperimentConfig]:
    """
    Test configurations for typical weight matrix shapes.
    Format: (hidden_dim, hidden_dim) for square matrices or (hidden_dim_in, hidden_dim_out)

    Note: Both M and N must be divisible by block_size (128)
    """
    # Common weight matrix shapes in transformers
    # Format: (in_features, out_features) for weight matrices
    dimensions = [1024, 2048, 4096, 8192, 16384]
    input_shapes = list(product(dimensions, repeat=2))

    configs = []
    block_sizes = [128]  # Standard block size for FP8

    for shape in input_shapes:
        for block_size in block_sizes:
            # Verify both dimensions are divisible by block_size
            if shape[0] % block_size == 0 and shape[1] % block_size == 0:
                configs.append(
                    ExperimentConfig(
                        input_shape=shape,
                        block_size=block_size,
                    )
                )
    return configs


def run_experiment(config: ExperimentConfig) -> ExperimentResult:
    """
    Run benchmark experiment comparing torch and Triton implementations.
    """
    M, N = config.input_shape
    block_size = config.block_size

    def torch_fp8_blockwise_weight_quant(
        x: torch.Tensor, block_size: int = 128
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Torch reference implementation for blockwise FP8 weight quantization.

        Quantizes in (block_size x block_size) blocks. Each block gets one scale factor.
        Outputs in column-major format for RHS operator.

        Args:
            x: Input tensor of shape (M, N), row-major
            block_size: Block size for quantization

        Returns:
            y: Quantized tensor in FP8, shape (M, N), column-major format
            s: Reciprocal scales in column-major format (M//block_size, N//block_size)
        """
        assert x.is_contiguous(), "Input must be contiguous"
        assert x.dim() == 2, "Input must be 2D"
        assert x.size(0) % block_size == 0 and x.size(1) % block_size == 0, (
            "Both dimensions must be divisible by block_size"
        )

        M, N = x.size()
        M_blocks = M // block_size
        N_blocks = N // block_size

        # FP8 E4M3 constants
        max_fp8_e4m3 = 448.0
        min_fp8_e4m3 = -448.0
        eps = 1e-12

        # Reshape to (M_blocks, block_size, N_blocks, block_size) for block-wise operations
        x_reshaped = x.view(M_blocks, block_size, N_blocks, block_size)
        # Permute to (M_blocks, N_blocks, block_size, block_size) for easier block processing
        x_blocks = x_reshaped.permute(0, 2, 1, 3)

        # Compute max absolute value per block (M_blocks, N_blocks)
        amax = torch.clamp(
            x_blocks.reshape(M_blocks, N_blocks, -1)
            .abs()
            .amax(dim=2)
            .to(torch.float64),
            min=eps,
            max=float("inf"),
        )

        # Compute scales (M_blocks, N_blocks)
        scale = (max_fp8_e4m3 / amax).to(torch.float32)

        # Broadcast scale for quantization (M_blocks, N_blocks, 1, 1)
        scale_broadcast = scale.unsqueeze(2).unsqueeze(3)

        # Quantize
        y_blocks = x_blocks * scale_broadcast
        y_blocks = torch.clamp(y_blocks, min=min_fp8_e4m3, max=max_fp8_e4m3)

        # Reshape back and convert to FP8
        # Permute back: (M_blocks, N_blocks, block_size, block_size) -> (M_blocks, block_size, N_blocks, block_size)
        y_reshaped = y_blocks.permute(0, 2, 1, 3)
        y_rowmajor = y_reshaped.reshape(M, N).to(torch.float8_e4m3fn)

        # Convert to column-major format

        y = y_rowmajor.t().contiguous().t()

        # Compute reciprocal scales - explicitly cast to float32
        reciprocal_scale = 1.0 / scale

        # Convert to column-major
        s = reciprocal_scale.t().contiguous().t()

        return y, s

    def verify_outputs(
        y_torch: torch.Tensor,
        s_torch: torch.Tensor,
        y_triton: torch.Tensor,
        s_triton: torch.Tensor,
        rtol: float = 1e-2,
        atol: float = 1e-2,
    ):
        """Verify that Triton and torch implementations produce similar results."""

        # Convert FP8 back to float for comparison

        y_torch_float = y_torch.to(torch.float32)
        y_triton_float = y_triton.to(torch.float32)

        assert y_torch.shape == y_triton.shape, (
            f"Output shape mismatch: torch {y_torch.shape} vs triton {y_triton.shape}"
        )
        assert y_torch.stride() == y_triton.stride(), (
            f"Output stride mismatch: torch {y_torch.stride()} vs triton {y_triton.stride()}"
        )

        assert s_torch.shape == s_triton.shape, (
            f"Scale shape mismatch: torch {s_torch.shape} vs triton {s_triton.shape}"
        )
        assert s_torch.stride() == s_triton.stride(), (
            f"Scale stride mismatch: torch {s_torch.stride()} vs triton {s_triton.stride()}"
        )

        # Check quantized values are close
        torch.testing.assert_close(
            y_torch_float,
            y_triton_float,
            rtol=rtol,
            atol=atol,
            msg="Quantized values differ between torch and Triton implementations",
        )

        torch.testing.assert_close(
            s_torch,
            s_triton,
            rtol=rtol,
            atol=atol,
            msg="Scales differ between torch and Triton implementations",
        )

    # Create input tensor
    input_tensor = torch.randn(
        M,
        N,
        dtype=torch.bfloat16,
        device=device,
    )

    # Benchmark torch implementation
    torch_impl_c = torch.compile(torch_fp8_blockwise_weight_quant)
    y_torch, s_torch = torch_impl_c(input_tensor, block_size)
    torch_time_us = benchmark_cuda_function_in_microseconds(
        torch_impl_c,
        input_tensor,
        block_size,
    )

    # Benchmark Triton implementation
    y_triton, s_triton = triton_fp8_blockwise_weight_quant_rhs(input_tensor, block_size)
    triton_time_us = benchmark_cuda_function_in_microseconds(
        triton_fp8_blockwise_weight_quant_rhs,
        input_tensor,
        block_size,
    )

    # Verify correctness
    verify_outputs(y_torch, s_torch, y_triton, s_triton)

    # Memory bandwidth calculations
    bytes_per_input_el = torch.finfo(input_tensor.dtype).bits / 8
    bytes_per_output_el = torch.finfo(y_triton.dtype).bits / 8
    bytes_per_scale_el = torch.finfo(s_triton.dtype).bits / 8

    read_bytes = input_tensor.numel() * bytes_per_input_el
    write_bytes = (
        y_triton.numel() * bytes_per_output_el + s_triton.numel() * bytes_per_scale_el
    )

    torch_gbps = ((read_bytes + write_bytes) / 1e9) / (torch_time_us / 1e6)
    triton_gbps = ((read_bytes + write_bytes) / 1e9) / (triton_time_us / 1e6)

    return ExperimentResult(
        torch_us=torch_time_us,
        triton_us=triton_time_us,
        torch_gbps=torch_gbps,
        triton_gbps=triton_gbps,
    )


def print_results(experiments: List[Experiment]):
    """Print benchmark results in a formatted table."""
    headers = [
        "input_shape (M, N)",
        "block_size",
        "torch_us",
        "triton_us",
        "speedup",
        "torch_gbps",
        "triton_gbps",
    ]
    rows = []
    for experiment in experiments:
        speedup = experiment.result.torch_us / experiment.result.triton_us
        rows.append(
            [
                f"{experiment.config.input_shape[0]}x{experiment.config.input_shape[1]}",
                experiment.config.block_size,
                f"{experiment.result.torch_us:.2f}",
                f"{experiment.result.triton_us:.2f}",
                f"{speedup:.2f}x",
                f"{experiment.result.torch_gbps:.1f}",
                f"{experiment.result.triton_gbps:.1f}",
            ]
        )
    print(tabulate(rows, headers=headers, tablefmt="grid"))


def main():
    """Main benchmark execution."""
    torch.random.manual_seed(123)
    configs = get_configs()
    results = []

    print(f"Running {len(configs)} benchmark configurations...\n")

    for config in tqdm(configs, desc="Benchmarking"):
        result = run_experiment(config)
        results.append(Experiment(config=config, result=result))

    print("\n" + "=" * 80)
    print("BENCHMARK RESULTS - RHS Weight Quantization")
    print("=" * 80 + "\n")
    print_results(results)


if __name__ == "__main__":
    main()
